import os
import random


def get_data_file(type):
    data_path = '/data_2/againstNoBelt_and_againstCalling/project/driver_detect/data/coco_' + type + '_data/'
    save_list = []
    img_dict = {}
    for file in os.listdir(data_path + 'images'):
        img_dict[file.strip()[:-4]] = 1

    for file in os.listdir(data_path + 'annotations'):
        if file[:-4] in img_dict:
            img_fullpath = os.path.join(data_path, 'images', file[:-4] + '.jpg')
            xml_fullpath = os.path.join(data_path, 'annotations', file)
            save_list.append(img_fullpath + ' ' + xml_fullpath)
    random.shuffle(save_list)
    print(len(save_list))
    with open(type + '.txt', 'w') as f:
        for line in save_list:
            f.write(line.strip() + '\n')


if __name__ == '__main__':
    get_data_file('train')
    get_data_file('test')